#include "odbc_database.h"
#include "exceptions.h"
#include "watch.h"
#include "sql.h"
#include "converter.h"
#include "str.h"

#define FETCH_FACTOR 4
#define DEFAULT_FETCH_SIZE 50 // fetch size; will be multiplied by FETCH_FACTOR each time that the size is increased

string odbc_database::get_error_details(SQLHANDLE handle, SQLSMALLINT type) const {
  SQLRETURN ret;

  SQLSMALLINT i = 1;
  SQLINTEGER native = 0;
  SQLCHAR state[7] = { "\0" };
  SQLCHAR text[256] = { "\0" };
  SQLSMALLINT len = 0;

  string result = global_empty_string;

  do {
    ret = SQLGetDiagRec(type, handle, i, state, &native, text, sizeof(text), &len);
    if (SQL_SUCCEEDED(ret)) {
      converter::to_string(i, str_cnv);

      result += str_cnv;
      result += ": ";
      result += (char*) state;
      result += ' ';
      result += (char*) text;
      result += "\r\n";

      ++i;
    } // if
  } while (ret == SQL_SUCCESS);

  return result;
} // odbc_database::get_error_details

void odbc_database::stop_transaction() {
  SQLRETURN ret = SQLSetConnectAttr(dbc, SQL_AUTOCOMMIT, (void*) SQL_AUTOCOMMIT_ON, 0);
  if (SQL_SUCCEEDED(ret) == false)
    throw db_exception("Cannot stop the transaction!");

  is_in_transaction = false;
} // odbc_database::stop_transaction

void odbc_database::process_columns(SQLHSTMT &stmt_handle, const unsigned int columns, record_set &rs) {
  SQLLEN indicator;
  SQLRETURN ret;
  buf[0] = L'\0';

  for (unsigned int j = 0; j < columns; ++j) {
    // retrieve column data as a string
    ret = SQLGetData(stmt_handle, j + 1, SQL_C_WCHAR, buf, sizeof(buf), &indicator); // column index is one-based here

    if (SQL_SUCCEEDED(ret) == false)
      break; // some error(s) occurred

    if (indicator == SQL_NULL_DATA)
      buf[0] = '\0'; // null column(s)

    converter::to_utf8(buf, str_cnv); // this does not work for binary data!
    rs.fill_record_set(str_cnv.c_str(), str_cnv.length()); // fills the record_set with data
  } // for j
} // odbc_database::process_columns

void odbc_database::connect(const string &connection_string, const string &protocol_file /*= global_empty_string*/) {
  watch w(true);

  if (connected == true)
    disconnect();

  protocol = protocol_file;

  SQLCHAR *dsn = (SQLCHAR*) connection_string.c_str();

  const SQLRETURN ret = SQLDriverConnect(dbc, 0, dsn, connection_string.length(), 0, 0, 0, SQL_DRIVER_COMPLETE);
  if (SQL_SUCCEEDED(ret) == false)
    throw db_exception("Could not connect to database!");

  connected = true;
  prot("CON", "odbc_database", w.get_elapsed_milliseconds());
} // odbc_database::connect

void odbc_database::connect(const string &host, const string &user_id, const string &password, const string &db_name, const string &protocol_file /*= global_empty_string*/) {
  string format = "DSN=%1;UID=%2;PWD=%3;";
  str::format(format, arg<string>(db_name, user_id, password)); // currently host will be ignored
  connect(format, protocol_file);
} // odbc_database::connect

void odbc_database::disconnect() {
  // close the connection to the database and cleanup
  if (connected == true) {
    watch w(true);

    rollback();
    SQLDisconnect(dbc);
    connected = false;

    prot("DCO", "odbc_database", w.get_elapsed_milliseconds());
  } // if
} // odbc_database::disconnect

bool odbc_database::is_connected() {
  if (connected == false)
    return false;

  SQLUINTEGER val = SQL_CD_TRUE;
  SQLRETURN ret = SQLGetConnectAttr(dbc, SQL_ATTR_CONNECTION_DEAD, &val, 0, 0);
  connected = (ret == SQL_SUCCESS && val == SQL_CD_FALSE); // SQL_CD_FALSE >> connection is OK
  return connected;
} // odbc_database::is_connected

void odbc_database::execute_query(const string &s, record_set &rs, bool append_to_record_set) {
  watch w(true);

  if (connected == false)
    throw db_exception(global_empty_string);

  wstring ws; // we have to convert the input back to UCS2 because of odbc api restrictions (UTF8 is not supported by odbc api)
  converter::to_ucs2(s, ws);
  SQLWCHAR *query = (SQLWCHAR*) ws.c_str();

  // allocate a statement handle
  SQLHSTMT stmt_handle;
  SQLRETURN ret = SQLAllocHandle(SQL_HANDLE_STMT, dbc, &stmt_handle);
  if (SQL_SUCCEEDED(ret) == false) {
    const string msg = get_error_details(stmt_handle, SQL_HANDLE_STMT);

    rollback();
    throw db_exception(msg);
  } // if

  ret = SQLExecDirectW(stmt_handle, query, s.length()); // third parameter can be SQL_NTS if the length is not known
  if (SQL_SUCCEEDED(ret) == false) {
    const string msg = get_error_details(stmt_handle, SQL_HANDLE_STMT);

    SQLFreeHandle(SQL_HANDLE_STMT, stmt_handle);
    stmt_handle = 0;
    rollback();
    throw db_exception(msg);
  } // if

  SQLSMALLINT columns = 0; // number of columns in result-set
  SQLSMALLINT column_name_length = 0;
  SQLLEN rows = 0; // number of rows in result-set
  SQLRowCount(stmt_handle, &rows);
  SQLNumResultCols(stmt_handle, &columns);

  const unsigned int n_rows = (unsigned int) rows;
  const unsigned int n_fields = (unsigned int) columns;

  if (append_to_record_set == false) {
    if (rows == -1)
      rs.reconstruct(DEFAULT_FETCH_SIZE, n_fields);
    else
      rs.reconstruct(n_rows, n_fields);

    for (unsigned int i = 0; i < n_fields; ++i) {
      SQLDescribeColW(stmt_handle, (SQLSMALLINT) (i + 1), buf, min(ODBC_BUFFER_SIZE, SHRT_MAX), &column_name_length, 0, 0, 0, 0);
      converter::to_utf8(buf, str_cnv);
      rs.fill_record_set(str_cnv.c_str(), column_name_length);
    } // for i
  } else if (rs.get_column_count() == n_fields) {
    if (rows == -1)
      rs.reserve(rs.get_row_count() + DEFAULT_FETCH_SIZE, n_fields);
    else
      rs.reserve(rs.get_row_count() + n_rows, n_fields);
  } else
    throw db_exception("Number of columns are not the same! (append_to_record_set = true)");

  if (rows == -1) { // number of rows not available, they must be fetched
    unsigned int row_counter = 0;
    unsigned int capacity = DEFAULT_FETCH_SIZE; // for rows

    while (SQL_SUCCEEDED(ret = SQLFetch(stmt_handle)) == true) {
      ++row_counter;

      if (capacity < row_counter) {
        capacity *= FETCH_FACTOR;
        rs.reserve(capacity, n_fields);
      } // if

      process_columns(stmt_handle, n_fields, rs);
    } // while

    rs.cut_rows_after(row_counter); // removes the last unused rows from record_set
  } else {
    for (unsigned int i = 0; i < n_rows; ++i) {
      if (SQL_SUCCEEDED(ret = SQLFetch(stmt_handle)) == false)
        break; // some error(s) occurred

      process_columns(stmt_handle, n_fields, rs);
    } // for i
  } // else

  SQLFreeHandle(SQL_HANDLE_STMT, stmt_handle);

  prot("QRY", s, w.get_elapsed_milliseconds(), rs.get_row_count());
} // odbc_database::execute_query

void odbc_database::execute_update(const string &s) {
  watch w(true);

  if (connected == false)
    throw db_exception(global_empty_string);

  wstring ws; // we have to convert the input back to UCS2 because of odbc api restrictions (UTF8 is not supported by odbc api)
  converter::to_ucs2(s, ws);
  SQLWCHAR *statement = (SQLWCHAR*) ws.c_str();

  // allocate a statement handle
  SQLHSTMT stmt_handle;
  SQLRETURN ret = SQLAllocHandle(SQL_HANDLE_STMT, dbc, &stmt_handle);
  if (SQL_SUCCEEDED(ret) == false) {
    const string msg = get_error_details(stmt_handle, SQL_HANDLE_STMT);

    rollback();
    throw db_exception(msg);
  } // if

  ret = SQLExecDirectW(stmt_handle, statement, s.length()); // third parameter can be SQL_NTS if the length is not known
  if (SQL_SUCCEEDED(ret) == false) {
    const string msg = get_error_details(stmt_handle, SQL_HANDLE_STMT);

    rollback();
    throw db_exception(msg);
  } // if

  ret = SQLFreeHandle(SQL_HANDLE_STMT, stmt_handle);
  if (SQL_SUCCEEDED(ret) == false) {
    rollback();
    throw db_exception(global_empty_string);
  } // if

  prot("UPD", s, w.get_elapsed_milliseconds());
} // odbc_database::execute_update

void odbc_database::start_transaction() {
  if (is_in_transaction == false) {
    watch w(true);

    const SQLRETURN ret = SQLSetConnectAttr(dbc, SQL_AUTOCOMMIT, (void*) SQL_AUTOCOMMIT_OFF, 0);
    if (SQL_SUCCEEDED(ret) == false)
      throw db_exception("Cannot start the transaction!");

    is_in_transaction = true;
    prot("TRA", "STARTED", w.get_elapsed_milliseconds());
  } // if
} // odbc_database::start_transaction

void odbc_database::commit() {
  if (is_in_transaction == true) {
    watch w(true);

    const SQLRETURN ret = SQLEndTran(SQL_HANDLE_DBC, dbc, SQL_COMMIT);
    if (SQL_SUCCEEDED(ret) == false)
      throw db_exception("Cannot commit the transaction!");

    stop_transaction();
    prot("TRA", "COMMIT", w.get_elapsed_milliseconds());
  } // if
} // odbc_database::commit

void odbc_database::rollback() {
  if (is_in_transaction == true) {
    watch w(true);

    const SQLRETURN ret = SQLEndTran(SQL_HANDLE_DBC, dbc, SQL_ROLLBACK);
    if (SQL_SUCCEEDED(ret) == false)
      throw db_exception("Cannot rollback the transaction!");

    stop_transaction();
    prot("TRA", "ROLLBACK", w.get_elapsed_milliseconds());
  } // if
} // odbc_database::rollback

void odbc_database::lock_table(const string &table_name) {
  // not supported
} // odbc_database::lock_table

void odbc_database::unlock_table(const string &table_name) {
  // not supported
} // odbc_database::unlock_table

odbc_database::odbc_database() : env(0), dbc(0) {
  buf[0] = L'\0';

  // allocate an environment handle
  if (SQLAllocHandle(SQL_HANDLE_ENV, SQL_NULL_HANDLE, &env) != SQL_SUCCESS)
    throw db_exception("Calling SQLAllocHandle failed!");

  // we want ODBC 3 support
  if (SQLSetEnvAttr(env, SQL_ATTR_ODBC_VERSION, (void*) SQL_OV_ODBC3, 0) != SQL_SUCCESS)
    throw db_exception("Calling SQLSetEnvAttr failed!");

  // allocate a connection handle
  if (SQLAllocHandle(SQL_HANDLE_DBC, env, &dbc) != SQL_SUCCESS)
    throw db_exception("Calling SQLAllocHandle failed!");

} // odbc_database::odbc_database

odbc_database::~odbc_database() {
  disconnect();

  SQLFreeHandle(SQL_HANDLE_DBC, dbc);
  SQLFreeHandle(SQL_HANDLE_ENV, env);

  dbc = env = 0;
} // odbc_database::~odbc_database
